import torch
from torch import Tensor
import iou3d_op
from utils import TicTocDecorator

def iou_core(box1: Tensor, box2: Tensor, area_sum: Tensor):
    overlap_w = torch.min(box1[2],box2[2]) - torch.max(box1[0],box2[0])
    overlap_h = torch.min(box1[3],box2[3]) - torch.max(box1[1],box2[1])
    if overlap_w <= 0 or overlap_h <= 0:
        return 0
    overlap_area = overlap_h * overlap_w
    return overlap_area / (area_sum - overlap_area)

@TicTocDecorator(10)
def iou_cpu(box1: Tensor, box2: Tensor):
    # box1 box2: [x1,y1,x2,y2]
    box1_width = box1[:,2]-box1[:,0]
    box1_height = box1[:,3]-box1[:,1]
    box1_area = box1_width*box1_height
    box2_width = box2[:,2]-box2[:,0]
    box2_height = box2[:,3]-box2[:,1]
    box2_area = box2_width*box2_height
    overlap_width = torch.clip(torch.min(box1[:,None,2],box2[None,:,2])-torch.max(box1[:,None,0],box2[None,:,0]),0.0)
    overlap_height = torch.clip(torch.min(box1[:,None,3],box2[None,:,3])-torch.max(box1[:,None,1],box2[None,:,1]),0.0)
    overlap = overlap_width*overlap_height
    iou = overlap/(box1_area[:,None]+box2_area[None,:]-overlap)
    return iou
@TicTocDecorator(10)
def iou_cpu2(box1: Tensor, box2: Tensor):
    box1_num = box1.size(0)
    box2_num = box2.size(0)
    box1_dim = box1.size(1)
    box2_dim = box2.size(1)
    if box1_dim != 4 or box2_dim != 4:
        return -1

    box1_area = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
    box2_area = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])

    inter_lr = torch.max(box1[:,None,:2],box2[None,:,:2])
    inter_rb = torch.min(box1[:,None,2:],box2[None,:,2:])
    inter = (inter_rb[...,0]-inter_lr[...,0]).clamp(min=0)*(inter_rb[:,:,1]-inter_lr[:,:,1]).clamp(min=0)

    union = box1_area[:,None] + box2_area[None,:] - inter

    iou = inter / union

    return iou

@TicTocDecorator(10)
def iou_cuda(box1: Tensor, box2: Tensor):
    result = torch.zeros(size=(len(box1), len(box2))).to(box1)
    flag = iou3d_op.forward_cuda(box1, box2, result)
    if flag == 0:
        return result
    return None
@TicTocDecorator(10)
def iou_cplus(box1: Tensor, box2: Tensor):
    result = torch.zeros(size=(len(box1), len(box2))).to(box1)
    flag = iou3d_op.forward_cplus(box1, box2, result)
    if flag == 0:
        return result
    return None
